/*
After bwa mapping, there is a need to filter those un-proper reads mapping,
like alignments with multiple hits, paired-end reads is inconsistent.

Distinguish proper paired-end mapping and unique single (including singleton) mapping with other not properly paired mapping and mapping with multiple hits
*/
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include "sam.h"

struct matrics {
	long pair; // proper paired-end
	long p_mr; // pair with multiple records
	long p_mh; // pair with multiple hits
	long p_no; // pair no hits
	long p_np; // pair not proper pair
	long p2s; // pair to singleton
	long p2s_mh; // singleton with multiple hits
	long single; // proper single-end
	long s_mh; // singled-end with multiple hits
	long s_no; // singled-end with no hits
};

#define is_mapped(b) (!((b)->core.flag & BAM_FUNMAP))
#define is_unique(b) ((bam_aux_get((b), "XT") != 0) && (bam_aux2A(bam_aux_get((b), "XT")) == 'U') && \
	(bam_aux_get((b), "XA") == 0))
#define is_proper_pair(a,b) (((a)->core.flag & BAM_FPROPER_PAIR) && \
	((b)->core.flag & BAM_FPROPER_PAIR) && \
	((a)->core.tid == (b)->core.tid) && \
	(((a)->core.flag & BAM_FREVERSE) != ((b)->core.flag & BAM_FREVERSE)))

int usage (void);
int readpair(samfile_t *, bam1_t **);
int checkpair(bam1_t **);
int checkalign(bam1_t *);

int main (int argc, char *argv[]){
	char in_mode[5], out_mode[5], *fn_matrics = 0, *fn_sgt = 0, *fn_log = 0;
	int is_bamin = 1, is_bamout = 1, with_header = 0, compress_level = -1, c = -1, flag = 0, not_end = 1;
	samfile_t *in = 0, *out = 0, *sgt = 0;
	FILE *LOG = stderr, *MATRICS = stderr;
	struct matrics m = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

	bam1_t * pair[2];

	strcpy(in_mode, "r"); strcpy(out_mode, "w");
	// handle commond-line arguments
	while((c = getopt(argc, argv, "?Suhm:s:l:")) >= 0){
		switch(c){
			case 'S': is_bamin = 0; break;
			case 'u': compress_level = 0; break;
			case 'h': with_header = 1; is_bamout = 0; break;
			case 'm': fn_matrics = strdup(optarg); break;
			case 's': fn_sgt = strdup(optarg); break;
			case 'l': fn_log = strdup(optarg); break;
			case '?':
			defualt: return usage();
		}
	}

	if(argc != optind + 2)	return usage(); // potential memory leak ...

	if(is_bamin)	strcat(in_mode, "b");
	if(is_bamout){
		strcat(out_mode, "b");
		if(compress_level >= 0){
			char tmp[2];
			tmp[0] = compress_level + '0'; tmp[1] = '\0';
			strcat(out_mode, tmp);
		}
	}else if(with_header)	strcat(out_mode, "h");
	// open BAM/SAM stream
	if((in = samopen(argv[optind], in_mode, 0)) == 0){
		fprintf(stderr, "[getUniqMapping] fail to open \"%s\" for reading.\n", argv[optind]);
		return 1;
	}
	optind ++;
	if((out = samopen(argv[optind], out_mode, in->header)) == 0){
		fprintf(stderr, "[getUniqMapping] fail to create \"%s\" for writing.\n", argv[optind]);
		return 1;
	}
	if(fn_sgt){
		if((sgt = samopen(fn_sgt, out_mode, in->header)) == 0){
			fprintf(stderr, "[getUniqMapping] fail to create \"%s\" for writing.\n", fn_sgt);
			return 1;
		}
	}else	sgt = out; // write singleton along with paired-end
	// open log file
	if(fn_log)
		if((LOG = fopen(fn_log, "w")) == 0){
			fprintf(stderr, "[getUniqMapping] fail to create log file \"%s\".\n", fn_log);
			LOG = stderr;
		}
	// init scan
	pair[0] = bam_init1(); pair[1] = bam_init1();
	not_end = 1;
	c = 0;

	long n = 0, j = 1;
	while(c = readpair(in, pair)){ // read a paired-end
		n += c;
		if(n >= j * 100000){
			fprintf(stderr, "[getUniqueMapping] processing %ld reads.\n", n);
			j ++;
		}
		if(c > 2){
			m.p_mr ++;
			fprintf(LOG, "%s p_mr\n", bam1_qname(pair[0]));
		}else if(c == 2){ // assume them two reads
			flag = checkpair(pair);
			if(flag & 16)	m.pair ++;
			else if(flag == 0){
				m.p_no ++;
				fprintf(LOG, "%s p_no\n", bam1_qname(pair[0]));
			}else if(flag == 1 || flag == 4){
				m.p2s ++;
				fprintf(LOG, "%s p2s\n", bam1_qname(pair[0]));
			}else if(flag == 2 || flag == 8){
				m.p2s_mh ++;
				fprintf(LOG, "%s p2s_mh\n", bam1_qname(pair[0]));
			}else if(flag == 10){
				m.p_mh ++;
				fprintf(LOG, "%s p_mh\n", bam1_qname(pair[0]));
			}else{
				m.p_np ++;
				fprintf(LOG, "%s p_np\n", bam1_qname(pair[0]));
			}

			if(flag & 16){
				samwrite(out, pair[0]);
				samwrite(out, pair[1]);
			}else{
				if(flag & 1)	samwrite(sgt, pair[0]);
				if(flag & 4)	samwrite(sgt, pair[1]);
			}
		}else if(c == 1){
			flag = checkalign(pair[0]);
			if(flag == 1)	m.single ++;
			else if(flag == 0){
				m.s_no ++;
				fprintf(LOG, "%s s_no\n", bam1_qname(pair[0]));
			}else if(flag){
				m.s_mh ++;
				fprintf(LOG, "%s s_mh\n", bam1_qname(pair[0]));
			}

			if(flag == 1)	samwrite(sgt, pair[0]);
		}
	}

	samclose(in);
	samclose(out);
	samclose(sgt);

	free(pair[0]); free(pair[1]);

	if(fn_matrics){
		if((MATRICS = fopen(fn_matrics, "w")) == 0){
			fprintf(stderr, "[getUniqMapping] fail to create file \"%s\" for writing.\n", fn_matrics);
			return 1;
		}
		fprintf(MATRICS, "##matrics\n");
		fprintf(MATRICS, "#pair: proper paired-end count\n");
		fprintf(MATRICS, "#p_mr: pair with multiple records\n");
		fprintf(MATRICS, "#p_mh: pair both with multiple hits\n");
		fprintf(MATRICS, "#p_no: pair both with no hits\n");
		fprintf(MATRICS, "#p_np: pair not proper paired\n");
		fprintf(MATRICS, "#p2s: pair with one no hits which another unique\n");
		fprintf(MATRICS, "#p2s_mh: pair with one no hits which another multiple hits\n");
		fprintf(MATRICS, "#single: single-end count\n");
		fprintf(MATRICS, "#s_mh: single-end with multiple hits\n");
		fprintf(MATRICS, "#s_no: single-end with no hit\n");
		fprintf(MATRICS, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", "pair", "p_mr", "p_mh", "p_no", "p_np", "p2s", "p2s_mh", "single", "s_mh", "s_no");
		fprintf(MATRICS, "%ld\t%ld\t%ld\t%ld\t%ld\t%ld\t%ld\t%ld\t%ld\t%ld\n", m.pair, m.p_mr, m.p_mh, m.p_no, m.p_np, m.p2s, m.p2s_mh, m.single, m.s_mh, m.s_no);
		fclose(MATRICS);
	}
	fclose(LOG);
	return 0;
}

/*
 * @para f: SAM/BAM stream
 * @para pair: paired-end alignment
 * @return: paired alignments count, while at more two of them update into pair
 */
int readpair (samfile_t *f, bam1_t *pair[2]){
	static char *qname = 0;
	static bam1_t *b = 0;
	register int c = 0, i = 0;

	bam_destroy1(pair[0]); bam_destroy1(pair[1]); // fixme: memory leak
	pair[0] = (bam1_t *)(0); pair[1] = (bam1_t *)(0);

	if(qname)	pair[i++] = bam_dup1(b);
	else{
		if(b == 0)	b = bam_init1();
		if((c = samread(f, b)) >= 0){
			qname = strdup(bam1_qname(b));
			pair[i++] = bam_dup1(b);
		}else{
			bam_destroy1(b);
			return 0;
		}
	}
	while((c = samread(f, b)) >= 0){
		if(strcmp(qname, bam1_qname(b)) == 0){
			if(i < 2){
				pair[i++] = bam_dup1(b);
			}else{
				i ++;
				while((c = samread(f,b)) >= 0 && strcmp(qname, bam1_qname(b)) == 0) i ++;
				break;
			}
		}else{
			free(qname); //fixme: memory leak
			qname = strdup(bam1_qname(b));
			break;
		}
	}
	if(c < 0){
		free(qname);
		qname = (char *)(0);
	}
	if(i == 0)	bam_destroy1(b);
	return i;
}

/*
 *@para p: two alignment paired
 *@return: a bit status which
	bit 1~2: first alignment is unmapped (0), unique (1) or multiple hits (2)
	bit 3~4: second alignment
	bit 5: proper pair (1) or not (0)
*/
int checkpair (bam1_t *p[2]){
	int flag = 0;
	flag |= checkalign(p[0]);
	flag |= ((checkalign(p[1]))<<2);
	if((flag == 5 || flag == 6 || flag == 9) && is_proper_pair(p[0], p[1]))	flag |= 16; // because bwa may set initiately unmapped reads as proper paired after WS local alignment (guess)
	return flag;
}
int checkalign (bam1_t *p){
	int flag = 0;
	if(is_mapped(p)){
		if(is_unique(p))	flag |= 1;
		else	flag |= 2;
	}
	return flag;
}

int usage (void){
	fprintf(stderr, "\n");
	fprintf(stderr, "Usage:    getUniqMapping [options] <in.bam>|<in.sam> <out.bam>\n\n");
	fprintf(stderr, "Options:  -S         Input is SAM. Default is BAM\n");
	fprintf(stderr, "          -u         Output is uncompressed BAM for pipe. Default is BAM, where all unique reads including paired-end or singleton reads are written to\n");
	fprintf(stderr, "          -h         Output is SAM with header\n");
	fprintf(stderr, "          -m file    Matrics file records count\n");
	fprintf(stderr, "          -s file    Output singleton unique reads into file.\n");
	fprintf(stderr, "          -l log     Log to log file\n");
	fprintf(stderr, "          -?         Show this help message\n\n");
	return 1;
}
